classdef CubeSolver < handle
    properties
        % physical parameter, problem independent
        lambda
        lengthD
        
        % physical parameter, problem dependent
        physPara      
        
        % nonphysical parameter
        nMax
        nGrid
        
        % generated
        unitGrid
        eigen
        matrixIdx
        
        % results
        Results
    end
    
    methods
        
        % initialize function
        function obj = CubeSolver(varargin)
            p = inputParser;
            
            % add nonphysical parameter
            p.addParameter('nGrid', 100, @(x) isscalar(x) && x>=10);
            p.addParameter('nMax', 10, @(x) x>1);
            
            % add physical parameter, problem independent
            p.addParameter('lambda', 0.0, @(x) x>=0.0);
            p.addParameter('lengthD', 1, @(x) x>0.0);
            
            % add physical parameter, problem dependent
            p.addParameter('tauD', 1.0, @(x) x>=0.0);
            p.addParameter('magD', 0.0, @(x) x>=0.0);
            p.addParameter('diffuseD', 0.0, @(x) x>0.0);
            p.addParameter('fS', 0.0, @(x) x>=0.0&&x<=1);
            p.addParameter('fF', 1.0, @(x) x>0.0&&x<=1);
            p.addParameter('OD', 1, @(x) x>0.0);
            p.addParameter('power', 1, @(x) x>0.0);
            p.addParameter('P0', 1, @(x) x>0.0);
            
            p.parse(varargin{:});
            
            obj.nGrid=p.Results.nGrid;
            obj.lengthD=p.Results.lengthD;
            obj.physPara=struct('tauD',p.Results.tauD, 'magD',p.Results.magD,...
                                'diffuseD',p.Results.diffuseD, 'fS',p.Results.fS,...
                                'fF',p.Results.fF, 'OD',p.Results.OD,...
                                'power',p.Results.power, 'P0',p.Results.P0);
                
            obj.setBoundaryCondition(p.Results.lambda, 'isUpdateEigen', false);
            obj.setMaxN(p.Results.nMax); % isUpdateEigen = true
            obj.setUnitGrid(p.Results.nGrid);
        end
        
        function obj = setBoundaryCondition(obj, lambda, varargin)
            p = inputParser;
            p.addRequired('lambda', @(x) isscalar(x) && x >= 0);
            p.addParameter('isUpdateEigen', true, @islogical);
            p.parse(lambda, varargin{:});
            
            obj.lambda = lambda;
            
            if p.Results.isUpdateEigen
                obj.updateEigen();
            end
        end
        
        function obj = setMaxN(obj, nMax, varargin)
            p = inputParser;
            p.addRequired('nMax', @(x) x>0);
            p.addParameter('isUpdateEigen', true, @islogical);
            p.parse(nMax, varargin{:});
            
            obj.nMax = p.Results.nMax;
            
            if p.Results.isUpdateEigen
                obj.updateEigen();
            end
        end
        
        function obj = setUnitGrid(obj, nGrid)
            p=inputParser;
            p.addRequired('nGrid', @(x) isscalar(x) && x>=10);
            p.parse(nGrid);
            
            obj.unitGrid.nGrid = 2*ceil(p.Results.nGrid/2); 
            obj.unitGrid.delta = 1.0/obj.unitGrid.nGrid;
            obj.unitGrid.posList = 0:obj.unitGrid.delta:1.0;
            [obj.unitGrid.X, obj.unitGrid.Y, obj.unitGrid.Z] = meshgrid(obj.unitGrid.posList, obj.unitGrid.posList, obj.unitGrid.posList);
        end
        
        function obj = setMatrixIdx(obj, nM)         
            vectors = { 1:nM,1:nM,1:nM};
            n = numel(vectors);
            combs = cell(1,n);
            [combs{end:-1:1}] = ndgrid(vectors{end:-1:1});
            combs = cat(n+1, combs{:});
            combs = reshape(combs,[],n);
            
            [~,sort_index]=sort(sum([obj.eigen.kList(combs(:,1)).^2,obj.eigen.kList(combs(:,2)).^2,obj.eigen.kList(combs(:,3)).^2],2));
            index=combs(sort_index,:);
            
            obj.matrixIdx.IdxX=index(:,1);
            obj.matrixIdx.IdxY=index(:,2);
            obj.matrixIdx.IdxZ=index(:,3);
            obj.matrixIdx.Idx3D=index;
        end
        
        function resM = overlapMatrix(obj, funcM, varargin)
            % calculate int(e_mnp*funM*e_ijk,cubic)
            p=inputParser; 
            p.addParameter('cutOff', obj.nMax, @(x) isscalar(x) && x>=1 && x<=obj.nMax );            
            p.parse(varargin{:});
            nM=p.Results.cutOff;
            obj.setMatrixIdx(nM);
            
            funcMVal=obj.getFunctionValue(funcM);
            
            resM=zeros(length(obj.matrixIdx.Idx3D),length(obj.matrixIdx.Idx3D));
            for i=1:length(obj.matrixIdx.Idx3D)
                for j=i:length(obj.matrixIdx.Idx3D)
                    e_mnp=obj.eigenMode(obj.matrixIdx.Idx3D(i,:));
                    e_ijk=obj.eigenMode(obj.matrixIdx.Idx3D(j,:));
                    resM(i,j)=simp3D(obj, funcMVal.*e_mnp.func3D.*e_ijk.func3D);
                    resM(j,i)=resM(i,j);
                end
            end
        end
        
        function resM = kappaMatrix(obj, varargin)
            % calculate k_mnp^2
            p=inputParser; 
            p.addParameter('cutOff', obj.nMax, @(x) isscalar(x) && x>=1 && x<=obj.nMax );            
            p.parse(varargin{:});
            nM=p.Results.cutOff;
            obj.setMatrixIdx(nM);
            
            resM = diag(obj.eigen.kList(obj.matrixIdx.IdxX).^2 + ...
                       obj.eigen.kList(obj.matrixIdx.IdxY).^2 + ...
                       obj.eigen.kList(obj.matrixIdx.IdxZ).^2);
        end
        
        function resM = projectMatrix(obj, funcM, varargin)
            % calculate int(e_mnp*funM,cubic)
            p=inputParser; 
            p.addParameter('cutOff', obj.nMax, @(x) isscalar(x) && x>=1 && x<=obj.nMax );            
            p.parse(varargin{:});
            nM=p.Results.cutOff;
            obj.setMatrixIdx(nM);
            
            funcMVal=obj.getFunctionValue(funcM);
            resM=zeros(length(obj.matrixIdx.Idx3D),1);
            for i=1:length(obj.matrixIdx.Idx3D)
                e_mnp=obj.eigenMode(obj.matrixIdx.Idx3D(i,:));
                resM(i)=simp3D(obj, funcMVal.*e_mnp.func3D);
            end
        end
        
        function SolveMz(obj, beamshape, varargin)
            p=inputParser; 
            p.addParameter('cutOff', obj.nMax, @(x) isscalar(x) && x>=1 && x<=obj.nMax );            
            p.parse(varargin{:});
            
            polarize=obj.lambertw3D(beamshape)./(1+obj.lambertw3D(beamshape));
            nu_z=(1+(3*polarize.^4+42*polarize.^2+35)./(3*polarize.^4+10*polarize.^2+3))/36;
            
            resM1 = obj.overlapMatrix(nu_z, 'cutOff', p.Results.cutOff);
            resM2 = obj.kappaMatrix('cutOff', p.Results.cutOff);
            
            funcVal = 0.5*obj.physPara.fS+obj.physPara.diffuseD*resM2+0.5*obj.physPara.fF*resM1;
            
            nu_mnp = obj.projectMatrix(0.25*obj.physPara.fS*polarize+0.25*obj.physPara.fF*polarize.*nu_z, 'cutOff', p.Results.cutOff);

            coeffMz = inv(funcVal)\nu_mnp;
            
            Mz=zeros(size(obj.unitGrid.X));
            for i=1:length(coeffMz)
                mode = obj.eigenMode(obj.matrixIdx.Idx3D(i,:));
                Mz =Mz + coeffMz(i)*mode.func3D;
            end
            obj.Results=struct('coeffMz',coeffMz,'Mz',Mz);
        end
        
        function SolveT2(obj, beamshape, vecVal, varargin)
            p=inputParser; 
            p.addParameter('cutOff', obj.nMax, @(x) isscalar(x) && x>=1 && x<=obj.nMax );            
            p.addParameter('tList',(0:0.1:1)',@iscolumn);
            p.parse(varargin{:});
            
            tList=p.Results.tList;
            polarize=obj.lambertw3D(beamshape)./(1+obj.lambertw3D(beamshape));
            nu_p=((polarize.^2+2).*(3*polarize.^4+90*polarize.^2+19)./(3*polarize.^4+10*polarize.^2+3))/36;
            
            resM1 = obj.overlapMatrix(nu_p, 'cutOff', p.Results.cutOff);
            resM2 = obj.kappaMatrix('cutOff', p.Results.cutOff);
            resM3 = obj.overlapMatrix(polarize, 'cutOff', p.Results.cutOff);
            
            funcVal = 0.5*obj.physPara.fS+obj.physPara.diffuseD*resM2+0.5*obj.physPara.fF*resM1+0.5*1i*obj.physPara.magD*resM3;
            
            eig_Mp = eig(funcVal);
            
            coeffMp=zeros(length(vecVal),length(tList));
            Mp=zeros(length(obj.unitGrid.posList),length(obj.unitGrid.posList),length(obj.unitGrid.posList),length(tList));
            for i = 1: length(tList)
                coeffMp(:,i)=expm(-funcVal*tList(i))*vecVal;
                for j=1:length(vecVal)
                    mode = obj.eigenMode(obj.matrixIdx.Idx3D(j,:));
                    Mp(:,:,:,i) = Mp(:,:,:,i) + coeffMp(j,i)*mode.func3D;
                end
            end
            obj.Results=struct('tList',tList,'eigMp',eig_Mp,'coeffMp',coeffMp,'Mp',Mp);
        end
        
        function res = lightPropagate(obj, beamshape)
            % coeff is transform constant between intensity and power: h*nu*L^2
            W = obj.lambertw3D(beamshape);
            
            res.P_in = simp3D(obj,repmat(obj.physPara.P0*W(:,:,1),1,1,length(obj.unitGrid.posList)));
            res.P_out = simp3D(obj,repmat(obj.physPara.P0*W(:,:,end),1,1,length(obj.unitGrid.posList)));
        end
        
        % disp functions for check
        function dispEigvalue(obj)
            kList = obj.eigen.kList(1:obj.nMax);
            xList=0:0.01:pi*obj.nMax;
            func=@(x) tan(x)-2*obj.lambda.*x./(x.^2-obj.lambda^2);
            yList=func(xList);
            yList(abs(yList)>10)=NaN;
            
            f = figure; ax=axes('Parent', f);
            plot(xList,yList,'r-');hold on;
            plot(kList,zeros(length(kList)),'b*');
            plot([min(xList),max(xList)],[0,0],'k--');
            legend('func','Solution')
            xlabel(ax, 'k'); 
        end
        
        function mode = eigenMode(obj, indexList, varargin)
            p=inputParser; p.KeepUnmatched = true;
            p.addRequired('indexList', @(x) length(x) == 3 && all(x>0) );            
            p.parse(indexList, varargin{:});
            
            kList = obj.eigen.kList(indexList); 
            deltaList = obj.eigen.deltaList(indexList);
            
            mode.kappa2 = kList'*kList;
            
            mode.funcX = sin(kList(1) * obj.unitGrid.posList + deltaList(1))./sqrt(0.5+obj.lambda/(kList(1)^2+obj.lambda^2));
            mode.funcY = sin(kList(2) * obj.unitGrid.posList + deltaList(2))./sqrt(0.5+obj.lambda/(kList(2)^2+obj.lambda^2));
            mode.funcZ = sin(kList(3) * obj.unitGrid.posList + deltaList(3))./sqrt(0.5+obj.lambda/(kList(3)^2+obj.lambda^2));
            mode.func3D = bsxfun(@times, mode.funcX.'*mode.funcY, reshape(mode.funcZ, 1, 1, []));
        end
        
        function dispMode(obj, indexList, varargin)
            p=inputParser; p.KeepUnmatched = true;
            p.addParameter('sliceX', 0.5, @(x) x>=0 && x<=1 );
            p.addParameter('sliceY', 0.5, @(x) x>=0 && x<=1 );
            p.addParameter('sliceZ', 0.5, @(x) x>=0 && x<=1 );
            p.parse(varargin{:});
            
            mode = obj.eigenMode(indexList, varargin{:});
            
            f = figure; ax=axes('Parent', f);
            slice(ax, obj.unitGrid.X*obj.lengthD, obj.unitGrid.Y*obj.lengthD, obj.unitGrid.Z*obj.lengthD, mode.func3D, ...
                  p.Results.sliceX*obj.lengthD, p.Results.sliceY*obj.lengthD, p.Results.sliceZ*obj.lengthD);
            c=colorbar(ax); c.Label.String = 'Polarization M_z';
            xlabel(ax, 'Position X (mm)'); ylabel(ax, 'Position Y (mm)'); zlabel(ax, 'Position Z (mm)');
            
        end
        
        function disp1DCheck(obj, varargin)
            % check solution of diffuseD*y''[x]-y[x]+1=0 with boundary
            % conditions lambda*y[0]-y'[0]=0,lambda*y[1]+y'[1]==0
            p=inputParser; 
            p.addParameter('cutOff', obj.nMax, @(x) isscalar(x) && x>=1 && x<=obj.nMax );            
            p.parse(varargin{:});
            
            cutOff=p.Results.cutOff;
            xList=0:0.01:1;xList=reshape(xList,1,[]);
            kList=obj.eigen.kList(1:cutOff);kList=reshape(kList,[],1);
            deltaList=obj.eigen.deltaList(1:cutOff);deltaList=reshape(deltaList,[],1);
            
            coeff=@(kx,deltax) (cos(deltax).*(1+(-1).^(obj.lambda>kx).*sign(cos(kx)))./...
                               (sqrt(0.5+obj.lambda./(kx.^2+obj.lambda^2)).*kx.*((obj.physPara.diffuseD)*kx.^2+1)));
            resCoeff=coeff(kList,deltaList);
            
            modeList=sin(kList*xList+deltaList) ./ sqrt(0.5+obj.lambda./(kList.*kList+obj.lambda*obj.lambda));
            exactFunc=@(x) 1-sqrt(obj.physPara.diffuseD)*(obj.lambda)*cosh((1-2*x)/(2*sqrt(obj.physPara.diffuseD)))./...
                              (sqrt(obj.physPara.diffuseD)*(obj.lambda)*cosh(0.5/sqrt(obj.physPara.diffuseD))+sinh(0.5/sqrt(obj.physPara.diffuseD)));
            
            exact=exactFunc(xList);
            appro=resCoeff'*modeList;
            f = figure; ax=axes('Parent', f);      
            plot(xList,exact);hold on;
            plot(xList,appro);
            legend(ax,{'exact',strcat('cutOff:',num2str(cutOff))});
            xlabel(ax, 'x');ylabel(ax, 'Mz(x)');
        end
        
        function dispPolarize(obj, beamshape)
            W=obj.lambertw3D(beamshape);
            P=W./(1+W);
            obj.plotCubic(P);
        end
    end
    
    %% private
    methods (Access =  private)
        
        function [kList, deltaList] = updateEigen(obj)
            if obj.lambda == 0
                kList = pi*(1:obj.nMax);
                deltaList = zeros(1, obj.nMax);
            elseif obj.lambda > 0
                tanPt = 0.5*pi*(1:2:2*obj.nMax); % divergent points of tan(x)
                intervalPt = sort([tanPt, obj.lambda]); % all divergent points, in ascent order
                
                f=@(x) tan(x)-2*obj.lambda.*x./(x.^2-obj.lambda^2);
                
                eps0 = 1e-3;% 1e-3 small enough ?
                eps1 = eps0 * min(abs(obj.lambda-tanPt)); 
                kList = arrayfun(@(k) fzero(f, [intervalPt(k)+eps1, intervalPt(k+1)-eps1]), 1:obj.nMax).';
                deltaList = atan(kList/obj.lambda);                
            else
                error('[lambda] must be nonnegative');
            end
            obj.eigen.kList = kList;
            obj.eigen.deltaList = deltaList;
        end
        
        function funcVal = getFunctionValue(obj, func)
            switch class(func)
                case 'function_handle'
                    funcVal = feval(func, obj.unitGrid.X, obj.unitGrid.Y, obj.unitGrid.Z);
                case 'double'
                    if all(size(func) == size(obj.unitGrid.X))
                        funcVal = func;
                    else
                        error('Function size does not match.');
                    end
                otherwise
                    error('Wrong type.');
            end
        end
    end
    
    %% external function
    methods
        out = simp3D(obj, U)
        res = lambertw3D(obj, beamshape)
        plotCubic(obj, funcM)
    end
end

